import numpy as np
import torch
import torch.nn.functional as F

from utils.train.imageUtils import to_categorical


def f_score(inputs, target, beta=1, smooth=1e-5, threhold=0.5):
    n, c, h, w = inputs.size()
    nt, ht, wt, ct = target.size()

    if h != ht and w != wt:
        inputs = F.interpolate(inputs, size=(ht, wt), mode="bilinear", align_corners=True)
    temp_inputs = torch.softmax(inputs.transpose(1, 2).transpose(2, 3).contiguous().view(n, -1, c), -1)
    temp_target = target.view(n, -1, ct)

    #   计算dice系数
    temp_inputs = torch.gt(temp_inputs, threhold).float()
    tp = torch.sum(temp_target * temp_inputs, axis=[0, 1])
    fp = torch.sum(temp_inputs, axis=[0, 1]) - tp
    fn = torch.sum(temp_target, axis=[0, 1]) - tp

    score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
    score = torch.mean(score)
    return score


def dsc(y_pred, y_true, num_classes):
    y_pred = np.round(y_pred).astype(int)
    y_true = np.round(y_true).astype(int)
    y_pred = to_categorical(y_pred, num_classes=num_classes)
    inSum = np.sum(y_pred * y_true)
    return inSum / (np.sum(y_pred) + np.sum(y_true) - inSum)
